package com.lft.tree09.self_balanceing_bst;

/**
 * 平衡二叉排序树
 */
public class AVLTreeDemo {
	public static void main(String[] args) {
	
	}
}

class AvlTree<E extends Comparable<E>> {
	private Node root;
	private int size;
	
	public AvlTree() {
		this.root = null;
		size = 0;
	}
	
	/**
	 * 获取某一节点的高度。
	 * @param node
	 * @return
	 */
	private int getHeight(Node node) {
		if (node == null) {
			return 0;
		}
		return node.height;
	}
	
	/**
	 * 获取元素个数。
	 * @return
	 */
	public int getSize() {
		return size;
	}
	
	/**
	 * 判断是否为空。
	 * @return
	 */
	public boolean isEmpty() {
		return size == 0;
	}
	
	/**
	 * 获取节点的平衡因子
	 * @param node
	 * @return
	 */
	private int getBalanceFactor(Node node) {
		if (node == null) {
			return 0;
		}
		return getHeight(node.left) - getHeight(node.right);
	}
	
	/**
	 * 判断树是否为平衡二叉树
	 * @return
	 */
	public boolean isBalanced() {
		return isBalanced(root);
	}
	
	private boolean isBalanced(Node node) {
		if (node == null) {
			return true;
		}
		int balanceFactory = Math.abs(getBalanceFactor(node));
		if (balanceFactory > 1) {
			return false;
		}
		return isBalanced(node.left) && isBalanced(node.right);
	}
	
	/**
	 * 右旋转
	 * @param y
	 * @return
	 */
	private Node rightRotate(Node y) {
		Node x = y.left;
		Node t3 = x.right;
		x.right = y;
		y.left = t3;
		
		// 更新 height
		y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
		x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
		return x;
	}
	
	/**
	 * 左旋转
	 * @param y
	 * @return
	 */
	private Node leftRotate(Node y) {
		Node x = y.right;
		Node t2 = x.left;
		x.left = y;
		y.right = t2;
		
		// 更新 height
		y.height = Math.max(getHeight(y.left), getHeight(y.right)) + 1;
		x.height = Math.max(getHeight(x.left), getHeight(x.right)) + 1;
		return x;
	}
	
	/**
	 * 向二分搜索树中添加新的元素(key,value)
	 * @param e
	 */
	public void add(E e) {
		root = add(root, e);
	}
	
	/**
	 * 向以 node 为根节点的二分搜索树中插入元素(key,value)，递归算法
	 * @param node
	 * @param e
	 * @return 返回插入新节点后二分搜索树的根节点。
	 */
	private Node add(Node node, E e) {
		if (node == null) {
			size++;
			return new Node(e);
		}
		// 要插入的节点小于当前节点
		if (e.compareTo(node.e) < 0) {
			// 向左递归查找要添加的位置。
			node.left = add(node.left, e);
		}
		// 要插入的节点大于当前节点
		else if (e.compareTo(node.e) > 0) {
			// 向右递归查找要添加的位置。
			node.right = add(node.right, e);
		}
		// 更新 height
		node.height = Math.max(getHeight(node.left), getHeight(node.right)) + 1;
		// 计算平衡因子
		int balanceFactor = getBalanceFactor(node);
		if (balanceFactor > 1 && getBalanceFactor(node.left) > 0) {
			// LL 右旋
			return rightRotate(node);
		}
		if (balanceFactor < -1 && getBalanceFactor(node.right) < 0) {
			// RR 左旋
			return leftRotate(node);
		}
		if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {
			// LR 先左旋左节点，再右旋
			node.left = leftRotate(node.left);
			return rightRotate(node);
		}
		if (balanceFactor < -1 && getBalanceFactor(node.right) > 0) {
			// RL 先右旋右节点，再左旋
			node.right = rightRotate(node.right);
			return leftRotate(node);
		}
		return node;
	}
	
	/**
	 * 删除元素
	 * @param e
	 * @return
	 */
	public E remove(E e) {
		Node node = getNode(root, e);
		if (node == null) {
			root = remove(root, e);
			return node.e;
		}
		return null;
	}
	
	/**
	 * 查找元素返回
	 * @param node
	 * @param e
	 * @return
	 */
	private Node getNode(Node node, E e) {
		if (node == null) {
			return null;
		}
		// 当前节点跟要查找的节点值相等，返回当前节点
		if (node.e.equals(e)) {
			return node;
		}
		if (e.compareTo(node.e) < 0) {
			// 如果小于当前节点值，递归查找左子树。
			return getNode(node.left, e);
		}
		if (e.compareTo(node.e) > 0) {
			// 如果大于当前节点值，递归查找右子树。
			return getNode(node.right, e);
		}
		return null;
	}
	
	private Node remove(Node node, E e) {
		if (node == null) {
			return null;
		}
		Node retNode = null;
		// 要删除的值比当前节点值小
		if (e.compareTo(node.e) < 0) {
			// 递归左子树。
			node.left = remove(node.left, e);
			retNode = node;
		}
		// 要删除的值比当前节点值大
		else if (e.compareTo(node.e) > 0) {
			// 递归右子树。
			node.right = remove(node.right, e);
			retNode = node;
		}
		// 要删除的值跟当前节点值相等。
		else {
			// 待删除的节点左子树为空。
			if (node.left == null) {
				Node rightNode = node.right;
				node.right = null;
				size--;
				retNode = rightNode;
			}
			// 待删除的节点右子树为空。
			else if (node.right == null) {
				Node leftNode = node.left;
				node.left = null;
				size--;
				retNode = leftNode;
			}
			// 待删除节点左右子树均不为空
			else {
				// 找到比待删除节点大的最小节点，即待删除节点右子树的最小节点
				Node successor = minNode(node.right);
				// 用这个节点顶替删除节点的位置
				successor.right = remove(node.right, successor.e);
				successor.left = node.left;
				node.left = node.right = null;
				retNode = successor;
			}
		}
		if (retNode == null) {
			return null;
		}
		// 维护平衡
		// 更新 height
		retNode.height = Math.max(getHeight(retNode.left), getHeight(retNode.right)) + 1;
		// 计算平衡因子
		int balanceFactor = getBalanceFactor(retNode);
		if (balanceFactor > 1 && getBalanceFactor(retNode.left) >= 0) {
			// LL 右旋
			return rightRotate(retNode);
		}
		if (balanceFactor < -1 && getBalanceFactor(retNode.right) <= 0) {
			// RR 右旋
			return leftRotate(retNode);
		}
		if (balanceFactor > 1 && getBalanceFactor(retNode.left) < 0) {
			// LR 先左旋，再右旋
			node.left = leftRotate(node.left);
			return rightRotate(retNode);
		}
		if (balanceFactor < -1 && getBalanceFactor(retNode.right) > 0) {
			// RL 先右旋，再左旋
			node.right = rightRotate(node.right);
			return leftRotate(retNode);
		}
		return retNode;
	}
	
	/**
	 * 找到当前树中的最小节点
	 * @param right
	 * @return
	 */
	private Node minNode(Node node) {
		if (node == null) {
			return node;
		}
		return minNode(node.left);
	}
	
	/**
	 * 内部节点类
	 */
	private class Node {
		public E e;
		public Node left;
		public Node right;
		public int height;
		
		public Node(E e) {
			this.e = e;
			this.left = null;
			this.right = null;
			this.height = 1;
		}
	}
}